from torchvision.datasets import MNIST
from python_ai.common.xcommon import *
import torch as pt
from torch.utils import data
import numpy as np

sep('MNIST mro')
print(MNIST.__mro__)

sep('my data set and loader')


class MyDataSet(data.Dataset):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.datat = pt.arange(1, 201)

    def __len__(self):
        return len(self.datat)

    def __getitem__(self, item):
        return self.datat[item]


my_data_set = MyDataSet()
my_data_loader = data.DataLoader(my_data_set,
                                 batch_size=32,
                                 shuffle=False,
                                 # drop_last=True,
                                 )

for i, bx in enumerate(my_data_loader):
    print(i, bx)
